from django.shortcuts import render, HttpResponse, redirect
from visualdl.server import app
from .static.my_code.predict import nn_infer
from .static.my_code.my_models import MyModels
from django.views.decorators.csrf import csrf_exempt
import json
from PIL import Image
from io import BytesIO
import matplotlib.pyplot as plt
from paddleseg.models import DeepLabV3, PSPNet, ANN, UNet, FCN, UNetPlusPlus
from paddleseg.models.backbones.resnet_vd import ResNet50_vd

# Create your views here.
def index(request):
    return render(request, 'index.html')

@csrf_exempt
def set_params(request):
    'ajax方式提交'
    # print(request.POST)
    netName = request.POST.get('netName')
    pretrainNet = request.POST.get('pretrainNet')
    num_classes = int(request.POST.get('numClasses'))
    image_size = int(request.POST.get('imageSize'))
    dataset_root = request.POST.get('rootDir')
    epochs = int(request.POST.get('epochs'))
    batch_size = int(request.POST.get('batchSize'))
    lr = float(request.POST.get('learningRate'))
    decay_steps = int(request.POST.get('decaySteps'))

    # my_model = MyModels(num_classes, image_size, dataset_root, epochs, batch_size, lr, decay_steps)
    # model = my_model.get_models(netName)
    # my_model.train_model(model, './')

    if netName == 'ANN':
        loss_img = 'img/loss/' + 'Train_loss_ann.png'
    elif netName == 'DeepLabV3':
        loss_img = 'img/loss/' + 'Train_loss_deeplabv3.png'
    elif netName == 'FCN':
        loss_img = 'img/loss/' + 'Train_loss_fcn.png'
    elif netName == 'PSPNet':
        loss_img = 'img/loss/' + 'Train_loss_pspnet.png'
    else:
        loss_img = 'img/loss/' + 'Train_loss_Unet.png'

    res = {'info': '参数设置完成！', 'loss_img': loss_img}
    return HttpResponse(json.dumps(res))

def get_params(request):
    pass

def show_net(request):
    app.run('./static/',
            model='net/model.pdmodel',
            host="127.0.0.1",
            port=8080,
            cache_timeout=20,
            language=None,
            public_path=None,
            api_only=False,
            open_browser=False)

    return redirect('http://127.0.0.1:8080/')

@csrf_exempt
def predict(request):
    # print(request.POST)
    # print(request.FILES)

    file_obj = request.FILES.get('originImg')
    f = open('index/static/img/upload/' + file_obj.name, mode='wb')
    for chunk in file_obj.chunks():
        f.write(chunk)
    f.close()

    img_path = 'index/static/img/upload/' + file_obj.name
    net_name = request.POST.get('netName')
    if net_name == 'ANN':
        model_params = 'index/static/model_params/ann/model.pdparams'
        model = ANN(num_classes=2, backbone=ResNet50_vd())
    elif net_name == 'DeepLabV3':
        model_params = 'index/static/model_params/deeplabv3/model.pdparams'
        model = DeepLabV3(num_classes=2, backbone=ResNet50_vd())
    elif net_name == 'FCN':
        model_params = 'index/static/model_params/fcn/model.pdparams'
        model = FCN(num_classes=2, backbone=ResNet50_vd())
    elif net_name == 'PSPNet':
        model_params = 'index/static/model_params/pspnet/model.pdparams'
        model = PSPNet(num_classes=2, backbone=ResNet50_vd())
    else:
        model_params = 'index/static/model_params/my_model/model.pdparams'
        model = UNet(num_classes=2)

    # image = Image.open(img_path)
    img_pred = nn_infer(model, img_path, model_params)
    plt.imshow(img_pred, cmap="gray")
    current_axes = plt.axes()
    current_axes.xaxis.set_visible(False)
    current_axes.yaxis.set_visible(False)
    plt.savefig('index/static/img/seg_res/' + file_obj.name, dpi=300, bbox_inches='tight')

    return HttpResponse(json.dumps({'info': '识别成功！', 'imgName': file_obj.name}))

@csrf_exempt
def get_res(request):
    file_type = request.GET.get('fileType')
    img_name = request.GET.get('imgName')
    if file_type == 'origin':
        img_path = 'index/static/img/upload/' + img_name
    else:
        img_path = 'index/static/img/seg_res/' + img_name

    img = Image.open(img_path)
    stream = BytesIO()
    img.save(stream, 'png')

    return HttpResponse(stream.getvalue())

@csrf_exempt
def download_img(request):
    img_name = request.GET.get('imgName')
    img_path = "index/static/img/seg_res/" + img_name
    with open(img_path, 'rb') as f:
        file = f.read()

    response = HttpResponse(file)
    response['Content-Type'] = 'application/octet-stream'
    response['Content-Disposition'] = 'attachment;filename={}'.format(img_name)
    return response